import matplotlib.pyplot as plt
import numpy as np

def visualize_small_grid(grid):
    size = len(grid)
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.set_xticks(np.arange(size + 1) - 0.5, minor=True)
    ax.set_yticks(np.arange(size + 1) - 0.5, minor=True)
    ax.grid(which="minor", color="black", linestyle='-', linewidth=0.5)
    ax.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
    
    # Display grid values
    cmap = plt.cm.get_cmap('viridis', 4)  # Map for values 0, 1, 2, 3
    img = ax.imshow(grid, cmap=cmap, origin='upper')
    plt.colorbar(img, ticks=[0, 1, 2, 3], label='Legend: 0 (Empty), 1 (Wall), 2 (Collector), 3 (Path)', fraction=0.046, pad=0.04)

    # Show values in each cell
    for i in range(size):
        for j in range(size):
            ax.text(j, i, str(grid[i][j]), ha='center', va='center', color='white' if grid[i][j] > 1 else 'black')

    plt.show()

# Example 44x44 grid with random data for 0, 1, 2, 3
np.random.seed(42)
example_grid = np.random.choice([0, 1, 2, 3], size=(44, 44), p=[0.7, 0.2, 0.05, 0.05])
visualize_small_grid(example_grid)
